{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "<a href=\"https://colab.research.google.com/drive/1PzWDDIItkrKbo2f8Reb20MqpafCR4VEc?usp=sharing\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"></a>"
      ],
      "metadata": {
        "id": "-xYqhOlfGOfn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Toolformer Agent Implementation\n",
        "\n",
        "A self-teaching agent that learns to use tools through:\n",
        "1. Generate potential API calls\n",
        "2. Execute actual API calls\n",
        "3. Observe real results\n",
        "4. Learn which tools improve predictions\n",
        "5. Self-fine-tune on successful tool usage"
      ],
      "metadata": {
        "id": "H7T7zJzjVE0a"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "ikWhmQQi0zn0"
      },
      "outputs": [],
      "source": [
        "!pip install -qU google-generativeai"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import google.generativeai as genai\n",
        "import getpass\n",
        "import re\n",
        "from datetime import datetime"
      ],
      "metadata": {
        "id": "KH7TgRdQVKeY"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Get Google's Gemini API Key here: https://aistudio.google.com/app/apikey"
      ],
      "metadata": {
        "id": "uaB23MaXXmCS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "API_KEY = getpass.getpass(\"Enter your Google API key: \")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UdM2O00yVMza",
        "outputId": "ef937880-e2f5-48ef-e9e6-8fa86beba1c5"
      },
      "execution_count": 3,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter your Google AI API key: ··········\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Configure API\n",
        "genai.configure(api_key=API_KEY)"
      ],
      "metadata": {
        "id": "j--pT9qsV_bL"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class ToolformerAgent:\n",
        "    def __init__(self):\n",
        "        self.model = genai.GenerativeModel(\"gemini-2.0-flash\")\n",
        "        self.tools = {}\n",
        "        self.tool_usage_history = []  # Track which tools work well\n",
        "        self.learned_patterns = []     # Successful tool usage patterns\n",
        "\n",
        "    def add_tool(self, name, func, description):\n",
        "        \"\"\"Register a tool with the agent\"\"\"\n",
        "        self.tools[name] = {\n",
        "            \"func\": func,\n",
        "            \"desc\": description,\n",
        "            \"success_count\": 0,\n",
        "            \"fail_count\": 0,\n",
        "            \"usefulness_score\": 0.5  # Start neutral\n",
        "        }\n",
        "\n",
        "    def generate_potential_calls(self, query):\n",
        "        \"\"\"Generate possible tool calls for the query\"\"\"\n",
        "        tools_desc = \"\\n\".join([\n",
        "            f\"- {name}: {info['desc']} (success rate: {self._get_success_rate(name)})\"\n",
        "            for name, info in self.tools.items()\n",
        "        ])\n",
        "\n",
        "        learned_examples = \"\\n\".join([\n",
        "            f\"Example: '{p['query']}' -> Use {p['tool']}({p['params']}) -> Success\"\n",
        "            for p in self.learned_patterns[-5:]  # Last 5 successful patterns\n",
        "        ]) if self.learned_patterns else \"No learned patterns yet.\"\n",
        "\n",
        "        prompt = f\"\"\"You are Toolformer, an agent that learns to use tools.\n",
        "\n",
        "        Query: {query}\n",
        "\n",
        "        Available Tools:\n",
        "        {tools_desc}\n",
        "\n",
        "        Learned Successful Patterns:\n",
        "        {learned_examples}\n",
        "\n",
        "        Generate 2-3 potential tool calls that might help answer this query.\n",
        "        For each, explain why it might be useful.\n",
        "\n",
        "        Format:\n",
        "        Tool: tool_name\n",
        "        Params: {{\"param\": \"value\"}}\n",
        "        Expected Help: How this tool would improve the answer\n",
        "        ---\n",
        "\n",
        "        Response:\"\"\"\n",
        "\n",
        "        response = self.model.generate_content(prompt).text\n",
        "        return self._parse_potential_calls(response)\n",
        "\n",
        "    def _parse_potential_calls(self, response):\n",
        "        \"\"\"Parse generated tool calls from LLM response\"\"\"\n",
        "        calls = []\n",
        "        current_call = {}\n",
        "\n",
        "        for line in response.split(\"\\n\"):\n",
        "            line = line.strip()\n",
        "            if line.startswith(\"Tool:\"):\n",
        "                if current_call:\n",
        "                    calls.append(current_call)\n",
        "                current_call = {\"tool\": line.split(\"Tool:\")[-1].strip()}\n",
        "            elif line.startswith(\"Params:\"):\n",
        "                try:\n",
        "                    params_str = line.split(\"Params:\")[-1].strip()\n",
        "                    current_call[\"params\"] = eval(params_str)\n",
        "                except:\n",
        "                    current_call[\"params\"] = {}\n",
        "            elif line.startswith(\"Expected Help:\"):\n",
        "                current_call[\"expected_help\"] = line.split(\"Expected Help:\")[-1].strip()\n",
        "            elif line == \"---\" and current_call:\n",
        "                calls.append(current_call)\n",
        "                current_call = {}\n",
        "\n",
        "        if current_call:\n",
        "            calls.append(current_call)\n",
        "\n",
        "        return calls\n",
        "\n",
        "    def execute_tool(self, tool_name, params):\n",
        "        \"\"\"Execute a tool and return result\"\"\"\n",
        "        try:\n",
        "            if tool_name not in self.tools:\n",
        "                return {\"success\": False, \"result\": None, \"error\": \"Tool not found\"}\n",
        "\n",
        "            result = self.tools[tool_name][\"func\"](**params)\n",
        "            return {\"success\": True, \"result\": result, \"error\": None}\n",
        "        except Exception as e:\n",
        "            return {\"success\": False, \"result\": None, \"error\": str(e)}\n",
        "\n",
        "    def evaluate_usefulness(self, query, tool_name, tool_result, original_answer, enhanced_answer):\n",
        "        \"\"\"Evaluate if tool improved the answer\"\"\"\n",
        "        prompt = f\"\"\"Evaluate if the tool improved the answer:\n",
        "\n",
        "        Query: {query}\n",
        "        Tool Used: {tool_name}\n",
        "        Tool Result: {tool_result}\n",
        "\n",
        "        Answer WITHOUT tool: {original_answer}\n",
        "        Answer WITH tool: {enhanced_answer}\n",
        "\n",
        "        Did the tool make the answer:\n",
        "        1. More accurate?\n",
        "        2. More informative?\n",
        "        3. More helpful?\n",
        "\n",
        "        Rate usefulness (0-10):\n",
        "        Reasoning:\"\"\"\n",
        "\n",
        "        response = self.model.generate_content(prompt).text\n",
        "\n",
        "        # Extract score\n",
        "        score_match = re.search(r'(\\d+)', response)\n",
        "        score = int(score_match.group(1)) if score_match else 5\n",
        "\n",
        "        return {\n",
        "            \"score\": score / 10.0,  # Normalize to 0-1\n",
        "            \"reasoning\": response,\n",
        "            \"improved\": score >= 7\n",
        "        }\n",
        "\n",
        "    def answer_without_tools(self, query):\n",
        "        \"\"\"Generate baseline answer without tools\"\"\"\n",
        "        prompt = f\"Answer this query briefly without using any tools: {query}\"\n",
        "        response = self.model.generate_content(prompt).text\n",
        "        return response.strip()\n",
        "\n",
        "    def answer_with_tool(self, query, tool_name, tool_result):\n",
        "        \"\"\"Generate enhanced answer using tool result\"\"\"\n",
        "        prompt = f\"\"\"Answer this query using the tool result:\n",
        "\n",
        "        Query: {query}\n",
        "        Tool Used: {tool_name}\n",
        "        Tool Result: {tool_result}\n",
        "\n",
        "        Provide an enhanced answer incorporating the tool result:\"\"\"\n",
        "\n",
        "        response = self.model.generate_content(prompt).text\n",
        "        return response.strip()\n",
        "\n",
        "    def learn_from_execution(self, query, tool_name, params, execution_result, evaluation):\n",
        "        \"\"\"Update learning based on tool execution outcome\"\"\"\n",
        "        tool_info = self.tools[tool_name]\n",
        "\n",
        "        if execution_result[\"success\"] and evaluation[\"improved\"]:\n",
        "            # Tool was useful\n",
        "            tool_info[\"success_count\"] += 1\n",
        "            tool_info[\"usefulness_score\"] = (\n",
        "                tool_info[\"usefulness_score\"] * 0.9 + evaluation[\"score\"] * 0.1\n",
        "            )\n",
        "\n",
        "            # Store successful pattern\n",
        "            self.learned_patterns.append({\n",
        "                \"query\": query,\n",
        "                \"tool\": tool_name,\n",
        "                \"params\": params,\n",
        "                \"score\": evaluation[\"score\"],\n",
        "                \"timestamp\": datetime.now().isoformat()\n",
        "            })\n",
        "\n",
        "            print(f\"✅ Learned: {tool_name} is useful for queries like '{query}'\\n\")\n",
        "\n",
        "        else:\n",
        "            # Tool wasn't helpful\n",
        "            tool_info[\"fail_count\"] += 1\n",
        "            tool_info[\"usefulness_score\"] = tool_info[\"usefulness_score\"] * 0.95\n",
        "            print(f\"❌ Learned: {tool_name} not useful for this query\\n\")\n",
        "\n",
        "        # Store usage history\n",
        "        self.tool_usage_history.append({\n",
        "            \"query\": query,\n",
        "            \"tool\": tool_name,\n",
        "            \"params\": params,\n",
        "            \"success\": execution_result[\"success\"],\n",
        "            \"useful\": evaluation[\"improved\"],\n",
        "            \"score\": evaluation[\"score\"],\n",
        "            \"timestamp\": datetime.now().isoformat()\n",
        "        })\n",
        "\n",
        "    def _get_success_rate(self, tool_name):\n",
        "        \"\"\"Calculate tool success rate\"\"\"\n",
        "        info = self.tools[tool_name]\n",
        "        total = info[\"success_count\"] + info[\"fail_count\"]\n",
        "        if total == 0:\n",
        "            return \"untested\"\n",
        "        return f\"{info['success_count']}/{total} ({info['usefulness_score']:.2f})\"\n",
        "\n",
        "    def run(self, query, auto_learn=True):\n",
        "        \"\"\"Run Toolformer with self-teaching\"\"\"\n",
        "        print(f\"\\n{'='*70}\")\n",
        "        print(f\"🎯 Query: {query}\")\n",
        "        print(f\"{'='*70}\\n\")\n",
        "\n",
        "        # Step 1: Generate baseline answer without tools\n",
        "        print(\"📝 Step 1: Generating baseline answer (no tools)...\")\n",
        "        baseline_answer = self.answer_without_tools(query)\n",
        "        print(f\"Baseline: {baseline_answer}\\n\")\n",
        "\n",
        "        # Step 2: Generate potential tool calls\n",
        "        print(\"🔧 Step 2: Generating potential tool calls...\")\n",
        "        potential_calls = self.generate_potential_calls(query)\n",
        "        print(f\"Generated {len(potential_calls)} potential tool calls:\\n\")\n",
        "\n",
        "        for i, call in enumerate(potential_calls, 1):\n",
        "            print(f\"  {i}. {call.get('tool', 'unknown')}({call.get('params', {})})\")\n",
        "            print(f\"     Why: {call.get('expected_help', 'N/A')}\\n\")\n",
        "\n",
        "        best_answer = baseline_answer\n",
        "        best_score = 0\n",
        "        best_tool = None\n",
        "\n",
        "        # Step 3 & 4: Execute each tool and observe results\n",
        "        for i, call in enumerate(potential_calls, 1):\n",
        "            tool_name = call.get(\"tool\")\n",
        "            params = call.get(\"params\", {})\n",
        "\n",
        "            if not tool_name or tool_name not in self.tools:\n",
        "                continue\n",
        "\n",
        "            print(f\"⚡ Step 3.{i}: Executing {tool_name}...\")\n",
        "            execution_result = self.execute_tool(tool_name, params)\n",
        "\n",
        "            if not execution_result[\"success\"]:\n",
        "                print(f\"   ❌ Execution failed: {execution_result['error']}\\n\")\n",
        "                continue\n",
        "\n",
        "            print(f\"   ✅ Result: {execution_result['result']}\\n\")\n",
        "\n",
        "            # Generate enhanced answer\n",
        "            print(f\"📊 Step 4.{i}: Evaluating usefulness...\")\n",
        "            enhanced_answer = self.answer_with_tool(query, tool_name, execution_result[\"result\"])\n",
        "\n",
        "            # Evaluate if tool improved answer\n",
        "            evaluation = self.evaluate_usefulness(\n",
        "                query, tool_name, execution_result[\"result\"],\n",
        "                baseline_answer, enhanced_answer\n",
        "            )\n",
        "\n",
        "            print(f\"   Score: {evaluation['score']:.2f}/1.0\")\n",
        "            print(f\"   Improved: {'Yes' if evaluation['improved'] else 'No'}\\n\")\n",
        "\n",
        "            # Step 5: Learn from execution\n",
        "            if auto_learn:\n",
        "                print(f\"🧠 Step 5.{i}: Learning from execution...\")\n",
        "                self.learn_from_execution(query, tool_name, params, execution_result, evaluation)\n",
        "\n",
        "            # Track best answer\n",
        "            if evaluation[\"score\"] > best_score:\n",
        "                best_score = evaluation[\"score\"]\n",
        "                best_answer = enhanced_answer\n",
        "                best_tool = tool_name\n",
        "\n",
        "        # Return best answer\n",
        "        print(f\"{'='*70}\")\n",
        "        print(f\"🏆 FINAL ANSWER\")\n",
        "        print(f\"{'='*70}\")\n",
        "        if best_tool:\n",
        "            print(f\"Best tool: {best_tool} (score: {best_score:.2f})\")\n",
        "        print(f\"\\n{best_answer}\\n\")\n",
        "\n",
        "        return best_answer\n",
        "\n",
        "    def show_learned_knowledge(self):\n",
        "        \"\"\"Display what the agent has learned\"\"\"\n",
        "        print(f\"\\n{'='*70}\")\n",
        "        print(\"🧠 LEARNED KNOWLEDGE\")\n",
        "        print(f\"{'='*70}\\n\")\n",
        "\n",
        "        print(\"Tool Performance:\")\n",
        "        for name, info in self.tools.items():\n",
        "            total = info[\"success_count\"] + info[\"fail_count\"]\n",
        "            if total > 0:\n",
        "                success_rate = info[\"success_count\"] / total * 100\n",
        "                print(f\"  {name}:\")\n",
        "                print(f\"    - Uses: {total}\")\n",
        "                print(f\"    - Success Rate: {success_rate:.1f}%\")\n",
        "                print(f\"    - Usefulness Score: {info['usefulness_score']:.2f}\")\n",
        "\n",
        "        print(f\"\\nSuccessful Patterns Learned: {len(self.learned_patterns)}\")\n",
        "        if self.learned_patterns:\n",
        "            print(\"\\nTop patterns:\")\n",
        "            sorted_patterns = sorted(self.learned_patterns, key=lambda x: x[\"score\"], reverse=True)\n",
        "            for pattern in sorted_patterns[:5]:\n",
        "                print(f\"  - Query: '{pattern['query']}'\")\n",
        "                print(f\"    Tool: {pattern['tool']}({pattern['params']})\")\n",
        "                print(f\"    Score: {pattern['score']:.2f}\\n\")"
      ],
      "metadata": {
        "id": "IcgfRKuIVQPI"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Define diverse tools\n",
        "def calculator(expression):\n",
        "    \"\"\"Perform mathematical calculations\"\"\"\n",
        "    try:\n",
        "        result = eval(expression)\n",
        "        return f\"Result: {result}\"\n",
        "    except Exception as e:\n",
        "        raise Exception(f\"Calculation error: {e}\")\n",
        "\n",
        "def search_web(query):\n",
        "    \"\"\"Search for current information (simulated)\"\"\"\n",
        "    # Simulated search results\n",
        "    results = {\n",
        "        \"weather\": \"Current weather: 22°C, Sunny\",\n",
        "        \"news\": \"Latest news: Tech conference announced for next month\",\n",
        "        \"python\": \"Python 3.12 released with performance improvements\",\n",
        "        \"ai\": \"New AI models showing improved reasoning capabilities\",\n",
        "    }\n",
        "\n",
        "    for key, val in results.items():\n",
        "        if key in query.lower():\n",
        "            return val\n",
        "\n",
        "    return f\"Search results for '{query}': General information available\"\n",
        "\n",
        "def calendar_check(date):\n",
        "    \"\"\"Check calendar for date (simulated)\"\"\"\n",
        "    schedules = {\n",
        "        \"today\": \"3 meetings scheduled: 10 AM, 2 PM, 4 PM\",\n",
        "        \"tomorrow\": \"1 meeting: 11 AM team sync\",\n",
        "        \"monday\": \"No meetings scheduled\",\n",
        "    }\n",
        "\n",
        "    return schedules.get(date.lower(), f\"No events found for {date}\")\n",
        "\n",
        "def translator(text, target_lang):\n",
        "    \"\"\"Translate text (simulated)\"\"\"\n",
        "    translations = {\n",
        "        \"hello\": {\"spanish\": \"Hola\", \"french\": \"Bonjour\", \"german\": \"Hallo\"},\n",
        "        \"goodbye\": {\"spanish\": \"Adiós\", \"french\": \"Au revoir\", \"german\": \"Auf Wiedersehen\"},\n",
        "    }\n",
        "\n",
        "    text_lower = text.lower()\n",
        "    if text_lower in translations and target_lang.lower() in translations[text_lower]:\n",
        "        return translations[text_lower][target_lang.lower()]\n",
        "\n",
        "    return f\"Translation: {text} -> {target_lang} (simulated)\"\n",
        "\n",
        "def unit_converter(value, from_unit, to_unit):\n",
        "    \"\"\"Convert between units\"\"\"\n",
        "    conversions = {\n",
        "        (\"km\", \"miles\"): lambda x: x * 0.621371,\n",
        "        (\"miles\", \"km\"): lambda x: x * 1.60934,\n",
        "        (\"kg\", \"lbs\"): lambda x: x * 2.20462,\n",
        "        (\"lbs\", \"kg\"): lambda x: x * 0.453592,\n",
        "        (\"celsius\", \"fahrenheit\"): lambda x: x * 9/5 + 32,\n",
        "        (\"fahrenheit\", \"celsius\"): lambda x: (x - 32) * 5/9,\n",
        "    }\n",
        "\n",
        "    key = (from_unit.lower(), to_unit.lower())\n",
        "    if key in conversions:\n",
        "        result = conversions[key](float(value))\n",
        "        return f\"{value} {from_unit} = {result:.2f} {to_unit}\"\n",
        "\n",
        "    raise Exception(f\"Conversion from {from_unit} to {to_unit} not supported\")\n"
      ],
      "metadata": {
        "id": "BF6MuhoJVYmd"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Usage Examples\n",
        "print(\"=\"*70)\n",
        "print(\"Toolformer Agentic Pattern : Self-Teaching Tool Usage\")\n",
        "print(\"=\"*70)\n",
        "\n",
        "agent = ToolformerAgent()\n",
        "\n",
        "# Register tools\n",
        "agent.add_tool(\"calculator\", calculator, \"Perform mathematical calculations\")\n",
        "agent.add_tool(\"search_web\", search_web, \"Search for current information\")\n",
        "agent.add_tool(\"calendar_check\", calendar_check, \"Check calendar and schedule\")\n",
        "agent.add_tool(\"translator\", translator, \"Translate text between languages\")\n",
        "agent.add_tool(\"unit_converter\", unit_converter, \"Convert between units\")\n",
        "\n",
        "# Example 1: Math query - should learn calculator is useful\n",
        "print(\"\\n\" + \"=\"*70)\n",
        "print(\"EXAMPLE 1: Math Query\")\n",
        "print(\"=\"*70)\n",
        "agent.run(\"What is 458 multiplied by 23 plus 1500?\")\n",
        "\n",
        "# Example 2: Current info - should learn search is useful\n",
        "print(\"\\n\" + \"=\"*70)\n",
        "print(\"EXAMPLE 2: Current Information Query\")\n",
        "print(\"=\"*70)\n",
        "agent.run(\"What's the latest news about AI?\")\n",
        "\n",
        "# Example 3: Scheduling - should learn calendar is useful\n",
        "print(\"\\n\" + \"=\"*70)\n",
        "print(\"EXAMPLE 3: Schedule Query\")\n",
        "print(\"=\"*70)\n",
        "agent.run(\"Do I have any meetings tomorrow?\")\n",
        "\n",
        "# Example 4: Translation - should learn translator is useful\n",
        "print(\"\\n\" + \"=\"*70)\n",
        "print(\"EXAMPLE 4: Translation Query\")\n",
        "print(\"=\"*70)\n",
        "agent.run(\"How do you say hello in Spanish?\")\n",
        "\n",
        "# Example 5: Unit conversion - should learn converter is useful\n",
        "print(\"\\n\" + \"=\"*70)\n",
        "print(\"EXAMPLE 5: Conversion Query\")\n",
        "print(\"=\"*70)\n",
        "agent.run(\"Convert 100 kilometers to miles\")\n",
        "\n",
        "# Example 6: Test learned knowledge - should now prefer successful tools\n",
        "print(\"\\n\" + \"=\"*70)\n",
        "print(\"EXAMPLE 6: Testing Learned Knowledge\")\n",
        "print(\"=\"*70)\n",
        "agent.run(\"Calculate 789 divided by 3\")\n",
        "\n",
        "# Show what the agent learned\n",
        "agent.show_learned_knowledge()"
      ],
      "metadata": {
        "id": "NW1ah_ZtVdc4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "z4_bBjVRWNkc"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}